from template import *
import random

numberDensity = 1.0e6
temperature = 1.5

me = 9.109382616e-31
#me = 29.0 * 1.67262158e-27


vte = math.sqrt(2.0*const.e*temperature/me)
vteVsim = math.sqrt(const.e*temperature/me)
vDrift = vte*1.0
#vDrift = 10000.0

coefFluxIne = 0.5 * (vte*math.exp(-vDrift*vDrift/(vte*vte))/math.sqrt(const.pi) + vDrift*(1.0+math.erf(vDrift/vte)))
particleFluxIne = numberDensity * coefFluxIne


coefFluxOute = 0.5 * (vte*math.exp(-vDrift*vDrift/(vte*vte))/math.sqrt(const.pi) - vDrift*(math.erfc(vDrift/vte)))
#coefFluxOute = 0.5 * (0.5*vte*math.exp(-vDrift*vDrift/(vte*vte)) - vDrift*(1.0-math.erf(vDrift/vte)))
#coefFluxOute = 0.5 * (0.5*vte*math.exp(-vDrift*vDrift/(vte*vte)) - vDrift*(1.0+math.erf(-vDrift/vte)))
particleFluxOute = numberDensity * coefFluxOute

print("vte: ", vte)
print("vDrift: ", vDrift)
print("")

print("particleFluxIne: ", particleFluxIne)
print("particleFluxOute: ", particleFluxOute)
print("")

m2kT = me / (2.0 * const.e * temperature)
m2pikT = me / (2.0 * const.pi * const.e * temperature)


numberV = 1000
minV = -5.0 * vte
maxV = 5.0 * vte
dV = (maxV - minV) / numberV

numberParticle = 100000

velocityList = np.zeros(numberV)
vdfTheoryList = np.zeros(numberV)
vdfTheoryInflowList = np.zeros(numberV)
vdfTheoryOutflowList = np.zeros(numberV)
vdfVsimInflowList = np.zeros(numberV)
vdfVsimOutflowList = np.zeros(numberV)

#================================= theory drifting maxwellian distrubution: begin =========================
coefFluxIneNumerical = 0.0
coefFluxOuteNumerical = 0.0
integration = 0.0
for i in range(0, numberV):
    velocityInMid = minV + 0.5 * dV + i * dV
    velocityList[i] = velocityInMid
    vdfTheoryList[i] = math.sqrt(m2pikT) * math.exp(-m2kT * (velocityInMid - vDrift) * (velocityInMid - vDrift))
    integration = integration + vdfTheoryList[i] * dV
    if velocityInMid < 0.0:
        coefFluxOuteNumerical = coefFluxOuteNumerical - velocityInMid * vdfTheoryList[i] * dV
    if velocityInMid > 0.0:
        coefFluxIneNumerical = coefFluxIneNumerical + velocityInMid * vdfTheoryList[i] * dV
print("integration: ", integration)
print("")
print("coefFluxIne: ", coefFluxIne)
print("coefFluxOute: ", coefFluxOute)
print("coefFluxOuteDiff: ", coefFluxIne-coefFluxOute)
print("")
print("coefFluxIneNumerical: ", coefFluxIneNumerical)
print("coefFluxOuteNumerical: ", coefFluxOuteNumerical)
print("coefFluxOuteNumericalDiff: ", coefFluxIneNumerical-coefFluxOuteNumerical)
print("")
#================================= theory drifting maxwellian distrubution: end =========================


#================================= theory inflow direction maxwellian distrubution: begin =========================
coefA = math.sqrt(m2pikT) / coefFluxIne
integration = 0.0
for i in range(0, numberV):
    velocityInMid = velocityList[i]
    vdfTheoryInflowList[i] = coefA * velocityInMid * math.exp(-m2kT * (velocityInMid - vDrift) * (velocityInMid - vDrift))
    if velocityInMid < 0.0:
        vdfTheoryInflowList[i] = 0.0
    integration = integration + vdfTheoryList[i] * dV
print("integration: ", integration)
print("")
#================================= theory inflow direction maxwellian distrubution: end =========================


#================================= theory outflow direction maxwellian distrubution: begin =========================
coefA = math.sqrt(m2pikT) / coefFluxOute
integration = 0.0
for i in range(0, numberV):
    velocityInMid = velocityList[i]
    vdfTheoryOutflowList[i] = -coefA * velocityInMid * math.exp(-m2kT * (velocityInMid - vDrift) * (velocityInMid - vDrift))  #/1.104438126803459
    if velocityInMid > 0.0:
        vdfTheoryOutflowList[i] = 0.0
    integration = integration + vdfTheoryOutflowList[i] * dV
print("integration: ", integration)
print("")
#================================= theory outflow direction maxwellian distrubution: end =========================


#================================= vsim inflow direction maxwellian distrubution: begin =========================
#vsim中算例等离子体从右向左流入
def EmitVxIne():
    global vteVsim
    global vDrift
    Vthe = vteVsim
    Vde = -vDrift

    i=0
    sigma2 = (1.+ (abs(Vde)/(2.*Vthe)))*Vthe**2
    alpha2 = 1 - (Vthe**2/sigma2)
    A = math.exp(abs(Vde)/Vthe)
    sigma=math.sqrt(sigma2)
    while i < 1:
        vx1=-sigma*math.sqrt(-2.0*math.log(random.random() ))
        zeta = random.random()*A*(vx1/Vthe**2)*math.exp(-vx1**2/(2.*sigma2))
        constraint = (vx1/Vthe**2)*math.exp(-(vx1-Vde)**2/(2.*Vthe**2))
        if zeta > constraint:
            i=i+1
            return vx1

dN = np.zeros(numberV)
for i in range(0, numberParticle):
    #print("i: ", i)
    velocitySample = -EmitVxIne()
    iVelocity = int((velocitySample - minV) / dV)
    if iVelocity < numberV:
        dN[iVelocity] = dN[iVelocity] + 1.0

integration = 0.0
for i in range(0, numberV):
    velocityInMid = velocityList[i]
    vdfVsimInflowList[i] = dN[i] / (numberParticle * dV)
    integration = integration + vdfVsimInflowList[i] * dV
print("integration: ", integration)
print("")
#================================= vsim inflow direction maxwellian distrubution: end =========================


#================================= vsim outflow direction maxwellian distrubution: begin =========================
#vsim中算例等离子体从右向左流入
def EmitVxOute():
    global vteVsim
    global vDrift
    Vthe = vteVsim
    Vde = vDrift

    i=0
    sigma2 = (1.+ (abs(Vde)/(2.*Vthe)))*Vthe**2
    alpha2 = 1 - (Vthe**2/sigma2)
    A = math.exp(abs(Vde)/Vthe)
    sigma=math.sqrt(sigma2)
    while i < 1:
        vx1=-sigma*math.sqrt(-2.0*math.log(random.random() ))
        zeta = random.random()*A*(vx1/Vthe**2)*math.exp(-vx1**2/(2.*sigma2))
        constraint = (vx1/Vthe**2)*math.exp(-(vx1-Vde)**2/(2.*Vthe**2))
        if zeta > constraint:
            i=i+1
            return vx1

dN = np.zeros(numberV)
for i in range(0, numberParticle):
    #print("i: ", i)
    velocitySample = EmitVxOute()
    iVelocity = int((velocitySample - minV) / dV)
    if iVelocity < numberV:
        dN[iVelocity] = dN[iVelocity] + 1.0

integration = 0.0
for i in range(0, numberV):
    velocityInMid = velocityList[i]
    vdfVsimOutflowList[i] = dN[i] / (numberParticle * dV)
    integration = integration + vdfVsimOutflowList[i] * dV
print("integration: ", integration)
print("")
#================================= vsim outflow direction maxwellian distrubution: end =========================

fig=plt.figure(figsize=(15,8))
	
fig.subplots_adjust(top=0.9,bottom=0.1,wspace=0.5,hspace=0.35)

ax=fig.add_subplot(1,1,1)

line0=ax.plot(velocityList,vdfTheoryList,lw=2, label = "theory")
line0=ax.plot(velocityList,vdfTheoryInflowList,lw=2, label = "inflow")
line0=ax.plot(velocityList,vdfTheoryOutflowList,lw=2, label = "outflow")
line0=ax.plot(velocityList,vdfVsimInflowList,lw=2, label = "inflow vsim")
line0=ax.plot(velocityList,vdfVsimOutflowList,lw=2, label = "outflow vsim")



ax.set_title("velocity distribution function")

#ax.set_xlim((0.0, 360.1))
#ax.set_xticks(np.arange(0, 361, step=60.0))
ax.set_xlabel("velocity", fontsize = label_fontsize)

ax.set_ylabel("vdf", fontsize = label_fontsize)

ax.legend()
ax.grid()

#figName = "vdf.png"
#fig.savefig(figName, dpi = 300)
plt.show()